{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "import IPython\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import soundfile as sf\n",
    "from tqdm import tqdm\n",
    "import tensorflow as tf\n",
    "\n",
    "from nara_wpe.tf_wpe import wpe\n",
    "from nara_wpe.utils import stft, istft, get_stft_center_frequencies\n",
    "from nara_wpe import project_root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stft_options = dict(size=512, shift=128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Minimal example with random data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def aquire_audio_data():\n",
    "    D, T = 4, 10000\n",
    "    y = np.random.normal(size=(D, T))\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "y = aquire_audio_data()\n",
    "Y = stft(y, **stft_options).transpose(2, 0, 1)\n",
    "with tf.Session() as session:\n",
    "    Y_tf = tf.placeholder(\n",
    "        tf.complex128, shape=(None, None, None))\n",
    "    Z_tf = wpe(Y_tf)\n",
    "    Z = session.run(Z_tf, {Y_tf: Y})\n",
    "z_tf = istft(Z.transpose(1, 2, 0), size=stft_options['size'], shift=stft_options['shift'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example with real audio recordings\n",
    "WPE estimates a filter to predict the current reverberation tail frame from K time frames which lie 3 (delay) time frames in the past. This frame (reverberation tail) is then subtracted from the observed signal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "channels = 8\n",
    "sampling_rate = 16000\n",
    "delay = 3\n",
    "iterations = 5\n",
    "taps = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Audio data\n",
    "Shape: (frames, channels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_template = 'AMI_WSJ20-Array1-{}_T10c0201.wav'\n",
    "signal_list = [\n",
    "    sf.read(str(project_root / 'data' / file_template.format(d + 1)))[0]\n",
    "    for d in range(channels)\n",
    "]\n",
    "y = np.stack(signal_list, axis=0)\n",
    "IPython.display.Audio(y[0], rate=sampling_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### STFT\n",
    "For simplicity reasons we calculate the STFT in Numpy and provide the result in form of the Tensorflow feed dict."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = stft(y, **stft_options).transpose(2, 0, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### iterative WPE\n",
    "A placeholder for Y is declared. The wpe function is fed with Y via the Tensorflow feed dict. Finally, an inverse STFT in Numpy is performed to obtain a dereverberated result in time domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nara_wpe.tf_wpe import get_power\n",
    "with tf.Session()as session:\n",
    "    Y_tf = tf.placeholder(tf.complex128, shape=(None, None, None))\n",
    "    Z_tf = wpe(Y_tf, taps=taps, iterations=iterations)\n",
    "    Z = session.run(Z_tf, {Y_tf: Y})\n",
    "z = istft(Z.transpose(1, 2, 0), size=stft_options['size'], shift=stft_options['shift'])\n",
    "IPython.display.Audio(z[0], rate=sampling_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Power spectrum\n",
    "Before and after applying WPE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, [ax1, ax2] = plt.subplots(1, 2, figsize=(20, 8))\n",
    "im1 = ax1.imshow(20 * np.log10(np.abs(Y[:, 0, 200:400])), origin='lower')\n",
    "ax1.set_xlabel('')\n",
    "_ = ax1.set_title('reverberated')\n",
    "im2 = ax2.imshow(20 * np.log10(np.abs(Z[:, 0, 200:400])), origin='lower')\n",
    "_ = ax2.set_title('dereverberated')\n",
    "cb = fig.colorbar(im1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
